#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
简化的大乐透模型训练脚本
"""

import os
import sys
import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import MinMaxScaler
import joblib

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, '..', '..'))
sys.path.append(project_root)

try:
    from model import LstmCRFModel
except ImportError as e:
    print(f"导入模型类失败: {e}")
    sys.exit(1)

class SimpleLotteryDataset(Dataset):
    def __init__(self, csv_file, window_size, red_balls=5, blue_balls=2):
        self.data = pd.read_csv(csv_file)
        self.scaler_X = MinMaxScaler()
        self.features, self.labels = self.preprocess(self.data, window_size, red_balls, blue_balls)

    def preprocess(self, data, window_size, red_balls, blue_balls):
        features, labels = [], []
        
        # 确保列名正确
        if '红球_1' in data.columns:
            # 重新命名列以匹配预期格式
            rename_dict = {}
            for i in range(1, red_balls + 1):
                rename_dict[f'红球_{i}'] = f'Red_{i}'
            for i in range(1, blue_balls + 1):
                rename_dict[f'蓝球_{i}'] = f'Blue_{i}'
            data.rename(columns=rename_dict, inplace=True)

        expected_columns = 1 + red_balls + blue_balls  # 期数 + 红球 + 蓝球
        if len(data.columns) < expected_columns:
            raise ValueError(f"数据列数不足，当前列数: {len(data.columns)}，期望至少 {expected_columns} 列。")

        for i in range(len(data) - window_size):
            # 特征：选取窗口内的红球和蓝球数据
            feature_window = data.iloc[i:i + window_size, 1:1 + red_balls + blue_balls].values
            features.append(feature_window)

            # 标签：下一期的红球和蓝球
            red_labels_seq = data.iloc[i + window_size, 1:1 + red_balls].values - 1  # 减1使其从0开始
            blue_label = data.iloc[i + window_size, 1 + red_balls:1 + red_balls + blue_balls].values - 1
            combined_labels = np.concatenate((red_labels_seq, blue_label))
            labels.append(combined_labels)

        # 转换为 NumPy 数组并进行缩放
        features_np = np.array(features)
        features_scaled = self.scaler_X.fit_transform(features_np.reshape(-1, features_np.shape[-1])).reshape(features_np.shape)
        labels_np = np.array(labels)

        return (
            torch.tensor(features_scaled, dtype=torch.float32),
            torch.tensor(labels_np, dtype=torch.long)
        )

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

def simple_train():
    # 配置
    DATA_FILE = os.path.join(current_dir, "dlt_history.csv")
    MODEL_PATH = os.path.join(current_dir, "dlt_lstm_model.pth")
    SCALER_PATH = os.path.join(current_dir, "scaler_X.pkl")
    WINDOW_SIZE = 5  # 减小窗口大小以适应较少的数据
    
    print("开始训练大乐透模型...")
    
    # 检查数据文件
    if not os.path.exists(DATA_FILE):
        print(f"错误: 数据文件不存在: {DATA_FILE}")
        return False
        
    # 加载数据
    print("加载数据...")
    try:
        dataset = SimpleLotteryDataset(DATA_FILE, WINDOW_SIZE)
        print(f"数据加载成功，共有 {len(dataset)} 个样本")
        
        if len(dataset) < 10:
            print("警告: 数据量较少，可能影响训练效果")
    except Exception as e:
        print(f"数据加载失败: {e}")
        return False
    
    # 创建数据加载器
    dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
    
    # 获取输入维度
    input_dim = dataset.features.shape[-1]
    print(f"输入维度: {input_dim}")
    
    # 检查 GPU 是否可用
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # 初始化模型
    red_model = LstmCRFModel(input_dim, hidden_dim=64, output_dim=35, output_seq_length=5, num_layers=1).to(device)
    blue_model = LstmCRFModel(input_dim, hidden_dim=64, output_dim=12, output_seq_length=2, num_layers=1).to(device)
    
    # 优化器
    red_optimizer = torch.optim.Adam(red_model.parameters(), lr=0.001)
    blue_optimizer = torch.optim.Adam(blue_model.parameters(), lr=0.001)
    
    # 训练几个epoch
    epochs = 50
    print(f"开始训练 {epochs} 个epoch...")
    
    for epoch in range(epochs):
        red_model.train()
        blue_model.train()
        
        total_red_loss = 0
        total_blue_loss = 0
        batch_count = 0
        
        for features, labels in dataloader:
            features = features.to(device)
            labels = labels.to(device)
            
            # 红球训练
            red_labels = labels[:, :5]
            red_mask = (red_labels >= 0)
            red_loss = red_model(features, red_labels, red_mask)
            red_optimizer.zero_grad()
            red_loss.backward()
            red_optimizer.step()
            total_red_loss += red_loss.item()
            
            # 蓝球训练
            blue_labels = labels[:, 5:]
            blue_mask = (blue_labels >= 0)
            blue_loss = blue_model(features, blue_labels, blue_mask)
            blue_optimizer.zero_grad()
            blue_loss.backward()
            blue_optimizer.step()
            total_blue_loss += blue_loss.item()
            
            batch_count += 1
        
        if (epoch + 1) % 10 == 0:
            avg_red_loss = total_red_loss / batch_count
            avg_blue_loss = total_blue_loss / batch_count
            print(f"Epoch {epoch+1}/{epochs}, 红球损失: {avg_red_loss:.4f}, 蓝球损失: {avg_blue_loss:.4f}")
    
    # 保存模型
    print("保存模型...")
    os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
    
    torch.save({
        "red_model": red_model.state_dict(),
        "blue_model": blue_model.state_dict()
    }, MODEL_PATH)
    
    # 保存缩放器
    joblib.dump(dataset.scaler_X, SCALER_PATH)
    
    print(f"模型已保存到 {MODEL_PATH}")
    print(f"缩放器已保存到 {SCALER_PATH}")
    print("训练完成！")
    
    return True

if __name__ == "__main__":
    success = simple_train()
    if success:
        print("大乐透模型训练成功！")
    else:
        print("大乐透模型训练失败！")